import os, sys
import numpy as np

class DCDS_Mask:
    def __init__(self):
        # masks for different number of pairs of sampples (reference/signal)
        self.npairs = [3, 4, 5, 6, 7, 8, 10]
        self.mask_3 = np.array([0, 0, 0, 0, 0, 0, 1, 1, 1, 0,
                                0, 0, 0, 0,-1,-1,-1, 0, 0, 0],
                                dtype=int)
        self.mask_4 = np.array([0, 0, 0, 0, 0, 1, 1, 1, 1, 0,
                                0, 0, 0,-1,-1,-1,-1, 0, 0, 0],
                                dtype=int)

        self.mask_5 = np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 0,
                                0, 0,-1,-1,-1,-1,-1, 0, 0, 0],
                                dtype=int)

        self.mask_6 = np.array([0, 0, 0, 0, 1, 1, 1, 1, 1, 1,
                                0,-1,-1,-1,-1,-1,-1, 0, 0, 0],
                                dtype=int)
        self.mask_7 = np.array([0, 0, 0, 1, 1, 1, 1, 1, 1, 1,
                                0,-1,-1,-1,-1,-1,-1,-1, 0, 0],
                                dtype=int)
        self.mask_8 = np.array([0, 0, 1, 1, 1, 1, 1, 1, 1, 1,
                                0,-1,-1,-1,-1,-1,-1,-1,-1, 0],
                                dtype=int)
        self.mask_10 = np.array([1, 1, 1, 1, 1, 1, 1, 1, 1, 1,
                                -1,-1,-1,-1,-1,-1,-1,-1,-1,-1],
                                dtype=int)

    def gen_mask(self, npairs):
        flag = npairs in self.npairs
        if flag is False:
            print('user provided npairs does not inclued in allowed range : {}'.format(self.npairs))
            sys.exit(0)

        if npairs == 3:
            return self.mask_3
        elif npairs == 4:
            return self.mask_4
        elif npairs == 5:
            return self.mask_5
        elif npairs == 6:
            return self.mask_6
        elif npairs == 7:
            return self.mask_7
        elif npairs == 8:
            return self.mask_8
        elif npairs == 9:
            return self.mask_9

###################################################################################################
if __name__ == '__main__':
    mask = DCDS_Mask()
    mask4 = mask.gen_mask(4)
    print('mask4 is {}'.format(mask4))

